import autograd
from .optimizer import Optimizer



class SGD(Optimizer):
    '''
    随机梯度下降
    '''

    def __init__(self, params, lr: float, momentum=0, weight_decay=0) -> None:
        if lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))

        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay)
        super().__init__(params, defaults)

    def step(self) -> None:
        for group in self.param_groups:
            weight_decay = group['weight_decay']
            momentum = group['momentum']
            lr = group['lr']

            for p in group['params']:
                if p.grad is None:
                    continue
                d_p = p.grad
                if weight_decay != 0:
                    d_p += weight_decay * p.data
                if momentum != 0:
                    pass
                p.data -= d_p * lr
